{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "729c2d52",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from shapely.geometry import LineString, Point\n",
    "import heapq\n",
    "\n",
    "# Parameters\n",
    "L = 1.0  # length of the arm\n",
    "pipe_radius = 0.1\n",
    "pipe_centers = [(1.0, 1.0), (1.5, 1.0), (2.0, 1.0)]\n",
    "\n",
    "# Workspace bounds\n",
    "x_vals = np.linspace(0, 3, 100)\n",
    "alpha_vals = np.linspace(0, np.pi/2, 100)\n",
    "\n",
    "# Create configuration space grid\n",
    "config_space = np.zeros((len(alpha_vals), len(x_vals)), dtype=bool)\n",
    "\n",
    "# Compute configuration space\n",
    "for i, alpha in enumerate(alpha_vals):\n",
    "    for j, x in enumerate(x_vals):\n",
    "        arm_start = (x, 0)\n",
    "        arm_end = (x + L * np.cos(alpha), L * np.sin(alpha))\n",
    "        arm_line = LineString([arm_start, arm_end])\n",
    "        collision = any(arm_line.distance(Point(px, py)) <= pipe_radius for (px, py) in pipe_centers)\n",
    "        config_space[i, j] = collision\n",
    "\n",
    "# Display configuration space\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.imshow(config_space, extent=[x_vals[0], x_vals[-1], np.degrees(alpha_vals[0]), np.degrees(alpha_vals[-1])],\n",
    "           origin='lower', cmap='Greys', aspect='auto')\n",
    "plt.xlabel('x (position)')\n",
    "plt.ylabel('alpha (degrees)')\n",
    "plt.title('Configuration Space (black = collision)')\n",
    "plt.grid(True)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ef0f95b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility functions for A* search\n",
    "def heuristic(a, b):\n",
    "    return np.linalg.norm(np.array(a) - np.array(b))\n",
    "\n",
    "def neighbors(pos, config_space):\n",
    "    i, j = pos\n",
    "    moves = [(-1, 0), (1, 0), (0, -1), (0, 1)]\n",
    "    result = []\n",
    "    for di, dj in moves:\n",
    "        ni, nj = i + di, j + dj\n",
    "        if 0 <= ni < config_space.shape[0] and 0 <= nj < config_space.shape[1]:\n",
    "            if not config_space[ni, nj]:  # Free space only\n",
    "                result.append((ni, nj))\n",
    "    return result\n",
    "\n",
    "def a_star(config_space, start, goal):\n",
    "    open_set = []\n",
    "    heapq.heappush(open_set, (0 + heuristic(start, goal), 0, start, [start]))\n",
    "\n",
    "    visited = set()\n",
    "\n",
    "    while open_set:\n",
    "        est_total_cost, cost_so_far, current, path = heapq.heappop(open_set)\n",
    "\n",
    "        if current in visited:\n",
    "            continue\n",
    "        visited.add(current)\n",
    "\n",
    "        if current == goal:\n",
    "            return path\n",
    "\n",
    "        for neighbor in neighbors(current, config_space):\n",
    "            if neighbor not in visited:\n",
    "                new_cost = cost_so_far + heuristic(current, neighbor)\n",
    "                est_total = new_cost + heuristic(neighbor, goal)\n",
    "                heapq.heappush(open_set, (est_total, new_cost, neighbor, path + [neighbor]))\n",
    "\n",
    "    return None  # No path found\n",
    "\n",
    "# Convert physical bin positions to indices in C-space grid\n",
    "def find_nearest_index(x_target, alpha_target_deg, x_vals, alpha_vals):\n",
    "    j = (np.abs(x_vals - x_target)).argmin()\n",
    "    i = (np.abs(np.degrees(alpha_vals) - alpha_target_deg)).argmin()\n",
    "    return (i, j)\n",
    "\n",
    "# Red bin -> start config, Blue bin -> goal config\n",
    "start_pos = find_nearest_index(0.9, 60, x_vals, alpha_vals)\n",
    "goal_pos = find_nearest_index(2.0, 60, x_vals, alpha_vals)\n",
    "\n",
    "# Run A* algorithm\n",
    "path = a_star(config_space, start_pos, goal_pos)\n",
    "\n",
    "# Plot path on configuration space\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.imshow(config_space, extent=[x_vals[0], x_vals[-1], np.degrees(alpha_vals[0]), np.degrees(alpha_vals[-1])],\n",
    "           origin='lower', cmap='Greys', aspect='auto')\n",
    "if path:\n",
    "    path_x = [x_vals[j] for i, j in path]\n",
    "    path_alpha = [np.degrees(alpha_vals[i]) for i, j in path]\n",
    "    plt.plot(path_x, path_alpha, 'r-', linewidth=2, label='Path')\n",
    "    plt.scatter(x_vals[start_pos[1]], np.degrees(alpha_vals[start_pos[0]]), color='green', label='Start')\n",
    "    plt.scatter(x_vals[goal_pos[1]], np.degrees(alpha_vals[goal_pos[0]]), color='blue', label='Goal')\n",
    "    plt.legend()\n",
    "else:\n",
    "    plt.title(\"No path found\")\n",
    "plt.xlabel('x (position)')\n",
    "plt.ylabel('alpha (degrees)')\n",
    "plt.title('Configuration Space with Path')\n",
    "plt.grid(True)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18e908f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "# Convert path indices to actual (x, alpha) values\n",
    "path_values = [(x_vals[j], alpha_vals[i]) for i, j in path]\n",
    "\n",
    "# Pipe drawing utility\n",
    "def draw_pipes(ax):\n",
    "    for (px, py) in pipe_centers:\n",
    "        pipe = plt.Circle((px, py), pipe_radius, color='gray')\n",
    "        ax.add_patch(pipe)\n",
    "\n",
    "# Animation function\n",
    "def animate_workspace_path(path_values, L):\n",
    "    fig, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "    def update(frame):\n",
    "        ax.clear()\n",
    "        ax.set_xlim(0, 3.5)\n",
    "        ax.set_ylim(0, 2)\n",
    "        ax.set_title('Robot Arm Workspace Animation')\n",
    "        ax.set_xlabel('x')\n",
    "        ax.set_ylabel('y')\n",
    "        ax.grid(True)\n",
    "\n",
    "        # Draw pipes\n",
    "        draw_pipes(ax)\n",
    "\n",
    "        # Draw red and blue bins\n",
    "        ax.add_patch(plt.Rectangle((0.85, 1.2), 0.1, 0.05, color='red'))\n",
    "        ax.add_patch(plt.Rectangle((1.95, 1.2), 0.1, 0.05, color='blue'))\n",
    "\n",
    "        # Get robot configuration\n",
    "        x, alpha = path_values[frame]\n",
    "        arm_start = (x, 0)\n",
    "        arm_end = (x + L * np.cos(alpha), L * np.sin(alpha))\n",
    "\n",
    "        # Draw arm\n",
    "        ax.plot([arm_start[0], arm_end[0]], [arm_start[1], arm_end[1]], 'b-', linewidth=4)\n",
    "        ax.plot(arm_start[0], arm_start[1], 'ko')  # base\n",
    "        ax.plot(arm_end[0], arm_end[1], 'ro')      # end effector\n",
    "\n",
    "    ani = animation.FuncAnimation(fig, update, frames=len(path_values), interval=200)\n",
    "    plt.close(fig)\n",
    "    return ani\n",
    "\n",
    "# Create the animation\n",
    "workspace_animation = animate_workspace_path(path_values, L)\n",
    "\n",
    "# Display in notebook\n",
    "# from IPython.display import HTML\n",
    "# HTML(workspace_animation.to_jshtml())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e8ca076",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the animation as a GIF\n",
    "gif_path = \"/mnt/data/robot_workspace_animation.gif\"\n",
    "workspace_animation.save(gif_path, writer='pillow', fps=5)\n",
    "gif_path\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
